#from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from pylab import *
import matplotlib.pyplot as plt
import time as timing #Renamed because time variable is already in use. Careful!
import math
import copy #Used to make a copy of grid objects.

import multiprocessing
from multiprocessing import Process, Pipe
#New instance either is all zeroes, or is copied off of an existing grid instance (template).
#Largely used as a storage unit for the displacement of the simulated sheet at different points.
#However, is also used to store mechanical properties of the sheet as values (this allows mechanical
#parameters to vary across the sheet, such as elastic modulus and density).
#In all honesty, this largely just creates functions to request pieces of the array
#in plain english. Grid.grid[:,0,1] vs Grid.getcolumn(0)
#Could be cleaned up, and sometimes Sheet just directly modifies the underlying
#numpy array. Bad form, I know...
class Grid:
	def __init__(self, xsize = 5, ysize = 5, tsize = 3, template = None):
		if template == None:
			self.grid = np.zeros((xsize, ysize, tsize))
			self.t = tsize
			self.xsize = xsize
			self.ysize = ysize
		else:
			self.grid = np.copy(template.grid)
			self.t = template.t		

	def getval(self, x, y, time=1):
		return self.grid[x,y,time]

	def setval(self, x, y, val, time=1):
		self.grid[x,y, time] = val

	def getrow(self, column, time=1):
		return self.grid[:,column, time]

	def getcolumn(self, row, time=1):
		return self.grid[row,:, time]

	def setcolumn(self, row, val, time=1):
		self.grid[row,:, time] = val
		
	def setall(self, val):
		for x in range(self.xsize):
			for y in range(self.ysize):
				if self.t == 1:
					self.grid[x,y] = val
				else:
					for t in range(self.t):
						self.grid[x,y,t] = val	
			
	def __eq__(self, othergrid):
		pass

	def multiply(self, val):
		self.grid = self.grid * val

	#Can directly set the grid class to match a numpy array of the same dimensions.
	def importnparray(self, nparray):
		self.grid = np.copy(nparray)

#Module for all the derivative calculations.
class MathModule:
	#Based off a recommendation of Paul Reimer.
	#Uses the symmetry assumption that near end points, a point distance x away from the
	#point in question should have the same value as the point a distance -x from the point.
	#Also, assumes a fixed end (displacement = 0 at boundary).
	def getpoints(self, point, distance, array, fixed = False):
		vals = []
		index = []
		start = point - distance
		end = point + distance + 1
		points = range(start, end)
		for val in points:
			if abs(val) < (array.size - 1):
				vals.append(val)
			else:
				if fixed:
					vals.append(None)
				else:
					vals.append((array.size - 1) - (val - (array.size - 1)))
		for a in vals:
			if a == None:
				index.append(0.0)
			else:
				index.append(array[a])			
		return index

	def secondderivative(self, array, point, stepsize):
		vals = self.getpoints(point, 1, array, fixed = False)		
		if self.equalvals(vals):
			return 0.0
		derivative = (vals[2] - 2*vals[1] + vals[0]) / stepsize**2
		return derivative

	def secondderivativehighaccuracy(self, array, point, stepsize):
		vals = self.getpoints(point, 2, array)
		if self.equalvals(vals):
			return 0.0

		derivative = (-vals[4] + 16*vals[3] - 30*vals[2] + 16*vals[1] - vals[0]) / 12*stepsize**2
		return derivative

	def fourthderivative(self, array, point, stepsize):
		vals = self.getpoints(point, 2, array, fixed = False)
		if self.equalvals(vals):
			return 0.0
		derivative = (vals[4] - 4*vals[3] + 6*vals[2] - 4*vals[1] + vals[0]) / stepsize**4	
		return derivative

	def fourthderivativehighaccuracy(self, array, point, stepsize):
		vals = self.getpoints(point, 3, array)
		if self.equalvals(vals):
			return 0.0
		derivative = (-vals[6] + 12*vals[5] - 39*vlas[4] + 56*vals[3] - 39*vals[2] + 12*vals[1] - vals[0]) / 6*stepsize**4
		return derivative

	def equalvals(self, vals):
		check = None
		for i in vals:
			if check == None:
				check = i
			else:
				if abs(check - i) >= 0.01:
					return False
		return True


#Represents the actual oscillating 2D sheet.
class Sheet:
	#Sets all internal variables for simulation.
	def __init__(self, xsize = 20, ysize = 20, timespacing=0.05, distspacing=0.5, setdrive='mid', damping=-0.2):
		self.xsize = xsize
		self.ysize = ysize
		self.xspacing = distspacing
		self.tspacing = timespacing
		self.grid = Grid(self.xsize, self.ysize)
		self.math = MathModule()
		self.setmaterialparameters()
		self.drivetype = setdrive
		self.dampingcoeff = damping

	#Sets all the 'physical' parameters, values that are mechanical properties of the sheet material.
	#User can either manually set values or program uses defaults.
	#Users can now 'unpickle' a numpy array for the emod and density, allowing for more complex
	#changes in mechanical parameters across the sheet.
	def setmaterialparameters(self):
		self.density = Grid(self.xsize, self.ysize, 1)
		self.elasticmodulus = Grid(self.xsize, self.ysize, 1)
		custom  = raw_input('Unpickle constants? (y/n): ')
		if custom == 'y':
			unpack()
		else:		
			modifymaterialparameters = bool(raw_input("Specify material parameters? (0 or 1): "))
			if not modifymaterialparameters:
				emod = 18.0	
				density = 4.0
				self.poisson = 0.37
			else:
				custom  = raw_input('Unpickle the elastic modulus? (y/n): ')
				if custom == 'y':
					unpack()
				else:
					emod = float(raw_input("Elastic modulus: "))
					density = float(raw_input("Density: " ))
				self.poisson = float(raw_input("Poisson ratio: "))

			self.elasticmodulus.setall(emod)
			self.density.setall(density)

		self.D = Grid(1,1,1,self.elasticmodulus)
		self.thickness=0.5
		const = self.thickness**3/(12*(1-self.poisson**2))
		self.D.multiply(const)	
		
		#Uses the pickle module to unpack a stored numpy array.	
		def unpack():
			import pickle
			print 'Pickled files must be of the Grid class!'
			emodpickle = raw_input("Name of pickled elastic modulus?: ")
			densitypickle = raw_input("Name of pickled density?: ")
			emodfile = open(emodpickle, 'rb')
			densityfile = open(densitypickle, 'rb')
			self.elasticmodulus.importnparray(pickle.load(emodfile))
			self.density.importnparray(pickle.load(densityfile))		

	#Interates through each timestep, rendering the sheet as it progresses.
	def mainloop(self, drivingfreq, endtime, maxamplitude, threads):
		print("Starting mainloop.")
		time = 0
		self.notified = False #See method checkstability().
		self.fixed = False
		self.amplitude = maxamplitude #Maximum amplitude of the sheet driver.
		self.renderscene()
		self.threadnumber = threads
		self.createchildren()
		print('Beginning calculations.')
		while time < endtime:
			try:
				self.drive(drivingfreq, self.tspacing, time, maxamplitude)
				self.fixedpoint(self.xsize/4, self.ysize/2)
				self.step()
				self.grid.grid[:,:,0] = np.copy(self.grid.grid[:,:,1])
				self.grid.grid[:,:,1] = np.copy(self.grid.grid[:,:,2])
				time += 1.0
				if time % 3 == 0:
					self.visualize2D(time)	
				print(int(time))	
				if self.notified == True:
					break
			except KeyboardInterrupt:
				print('\n Interrupted, stopping')
				break

		response = raw_input('Close? (Hit Enter) ')	
		plt.close('all')
		
		for child in self.childarray:
			child.terminate()
	
	#Makes a point on the sheet remain unmoving.
	def fixedpoint(self,x,y):
		self.grid.setval(x,y,0)
		self.grid.setval(x,y-1,0)
		self.grid.setval(x,y+1,0)
		self.grid.setval(x+1,y,0)
		self.grid.setval(x+1,y-1,0)
		self.grid.setval(x+1,y+1,0)
	#Creates the inital scene.
	def renderscene(self):
		self.X = np.arange(0, self.xsize)
		self.Y = np.arange(0, self.ysize)
		self.X, self.Y = np.meshgrid(self.X, self.Y)
		self.Z = self.grid.getval(self.X, self.Y)
		self.fig = plt.figure(figsize = (14, 11), dpi = 80)
		self.ax = self.fig.gca(projection='3d')
		self.surf = self.ax.plot_surface(self.X, self.Y, self.Z, rstride=1, cstride=1, cmap=cm.coolwarm, linewidth=0, antialiased=False)

	#Pause/render method for visualizing wave propogation.
	def visualize2D(self, settime, done=False):
		self.surf.remove()
		self.Z = self.grid.getval(self.X, self.Y)
		self.surf = self.ax.plot_surface(self.X, self.Y, self.Z, rstride=1, cstride=1, cmap=cm.jet, vmin = -self.amplitude*1.3, vmax = self.amplitude*1.3, linewidth=0, antialiased=True)
		self.ax.set_zlim((-int(10),int(10)))
		if settime % 2 == 0:
			plt.draw()
		if not done:
			val = int(settime)
			if val < 10:
				val = '000' + str(val)
			elif val < 100:
				val = '00' + str(val)
			elif val < 1000:
				val = '0' + str(val)
			elif val < 10000:
				val = str(val)
			fname = 'image' + str(self.drivetype) + str(val) + '.png'
			#Uncomment the next line if making an animation. 			
			#matplotlib.pyplot.savefig(fname, edgecolo='w', format='png')
		timing.sleep(0.01)		


	#Creates the effect of a driving force on the sheet, either at a point or in a line. 
	def drive(self, drivingfreq, timespacing, curtime, amplitude):
		sin, pi = np.sin, np.pi
		oscillation = amplitude * sin(np.pi/2 * drivingfreq * timespacing * curtime)
		#if curtime > 10 and abs(oscillation) < 0.05 and self.fixed == False:
		#	oscillation = 0
		#	self.fixed = True
		#	print('Fixed')
		#if self.fixed:
		#	oscillation = 0
		#oscillation = amplitude
		#self.grid.setcolumn(0, oscillation)
		if self.drivetype == 'mid':
			self.grid.setval(self.xsize/2, self.ysize/2,oscillation)
		if self.drivetype == 'side':
			self.grid.setcolumn(0, oscillation)
		if self.drivetype == 'end':
			self.grid.setval(self.xsize/2, 0, oscillation)

	#Checks if a result value on sheet exceeds the max driving frequency amplitude.
	#If the warning is thrown, this probably means the amplitude is growing  unbounded (unrealistic).
	def checkstability(self, val, amplitude, x, y):
		if abs(val) >= 10.0 * amplitude and not self.notified:
			print('Instability!')
			self.notified = True
	
	#Multithreading!
	#Creates each child process and associated pipe for giving tasks.
	def createchildren(self):
		self.childarray = []
		self.pipearray = []
		self.sizeperthread = self.xsize / self.threadnumber
		childID = 0
		print("Creating children")
		for x in range(self.threadnumber):
			parentpipe, childpipe = Pipe()
			newchild = Process(target=self.shake(childpipe, childID))
			newchild.start()
			self.childarray.append(newchild)
			self.pipearray.append(parentpipe)
			childID += 1
		
	#Breaks down the sheet into even chunks along the x axis for the child processes.
	#Then reassembles the grid based on the returns from the children.
	def step(self):
		start = 0
		for x in self.pipearray:
			x.send([start, start+self.sizeperthread,np.copy(self.grid.grid)])
			start += self.sizeperthread

		gridarray = []
		for pipe in self.pipearray:
			val = pipe.rcv()
			gridarray.append(val)

		increment = -1
		val = 0	
		newarray = None
		while increment != self.threadnumber:
			for x in gridarray:
				if x[0] == increment + 1:
					if newarray == None:
						newarray = np.copy(x[1])
					else:
						newarray = np.concatenate((newarray, x[1]))
			increment += 1
		
		self.grid.grid[:,:,2] = np.copy(newarray)
		return 0
		
	#Called function by step().
	#Used to actually simulate the shaking of the sheet.
	#Because child processes inherit memory space, we don't need args!...besides the duplex pipe, child ID,
	#and grid portion that is the child's responsibility.
	#Puts the calculated portion of the grid into the pipe.
	def shake(self,pipe,childID):
		print('Child ' + str(childID) + ' starting!')
		while True:
			if pipe.poll():
				print('Results')	
				vals = pipe.recv()
				portion = vals[:2]
				self.grid.grid = vals[2]
				for x in range(portion[0],(portion[1])):
					column = self.grid.getcolumn(x)
					dx2column = self.constructdx2(x)
					for y in range(self.ysize):
						row = self.grid.getrow(y)
						increment = self.increment(x, y, dx2column, row, column)
						self.checkstability(increment, self.amplitude, x, y)
						self.grid.setval(x, y, increment, 2)

				gridportion = np.copy(self.grid.grid[portion[0]:portion[1],:,2])
				pipe.send([childID, gridportion])
			else:
				timing.sleep(0.01)

	#Finds d2w/dx2, for use to find d4w/dx2dy2.
	def constructdx2(self, x):
		dwdx2column = np.zeros((self.ysize))
		for y in range(self.ysize):
			row = self.grid.getrow(y)
			dwdx2 = self.math.secondderivative(row, x, self.xspacing) 
			dwdx2column[y] =  dwdx2
		return dwdx2column	

	#Individual point change calculation, applied at every point on grid.
	def increment(self, x, y, dx2array, row, column):
		dwdy4 = self.math.fourthderivative(column, y, self.xspacing)
		dwdx4 = self.math.fourthderivative(row, x, self.xspacing)				
		dwdx2dy2 = self.math.secondderivative(dx2array[:], y, self.xspacing)
		seg1 = self.D.getval(x,y,0)*(dwdy4 + dwdx4 + dwdx2dy2)
		seg2 =  self.getdensity(x,y) * self.thickness*(-2*self.grid.getval(x,y)+self.grid.getval(x,y,0))/self.tspacing**2
		seg3 =  self.damping(x,y)*-self.grid.getval(x,y,0)/(2*self.tspacing)
		seg4 = 	self.getdensity(x,y)*self.thickness/self.tspacing**2+self.damping(x,y)/(2*self.tspacing)
		increment = -(seg1 + seg2 + seg3)/seg4
		self.checkstability(increment, self.amplitude, x, y)
		return increment	

	#Calculates the damping for a position on the sheet (should be greatest near the edge).
	def damping(self, x, y):
		coeff = 0.0
		damping_thickness_x = 0.15
		damping_thickness_y = 0.05
		if y <= self.ysize * damping_thickness_y or y >= self.ysize * (1.0 - damping_thickness_y):
			coeff = self.dampingcoeff / (self.ysize * 0.1) * y * 0.0
		if x >= self.xsize * (1.0 - damping_thickness_x):
			coeff = self.dampingcoeff / (self.xsize * 0.8) * x
		return coeff

	def getdensity(self, x, y):
		val = 1.0
		return val

#Quickly sets the components necessary for a basic simulation.
def sim(xsize=15, ysize = 15, time=20, drive='side'):
	threads = thread(xsize)
	tspacing = 0.05 
	xspacing = 0.5
	matplotlib.interactive(True)
	drivingamplitude = 1.0
	drivingfreq = 0.5
	damping = 0.2
	sheet = Sheet(xsize, ysize, tspacing, xspacing, setdrive=drive, damping=damping)
	sheet.mainloop(drivingfreq, time, drivingamplitude, threads)	 


def thread(xsize):
	extra = 1
	while extra != 0:
		threads = int(raw_input('Number of threads? (must be a divisor of ' + str(xsize) + '): '))
		extra = xsize % threads
	return threads 
